from typing import Optional, Union
import asyncio
import random

from axiom_boot.conf.manager import Settings
from axiom_boot.di import service, order, conditional_on_setting
from axiom_boot.logging.setup import get_logger
from axiom_boot.scraper.interfaces import ScraperMiddleware
from axiom_boot.scraper.models import Target, Response


logger = get_logger(__name__)


@order(20)  # 在日志(10)之后，代理(50)和重试(100)之前执行
@service()
@conditional_on_setting(key="scraper__rate_limit_enabled", expected_value=True)
class RateLimitingMiddleware(ScraperMiddleware):
    """
    一个爬虫中间件，用于通过限制并发请求数和增加请求间延时来控制爬取频率。
    它使用 asyncio.Semaphore 限制并发数，并使用用户配置的延时范围为每个请求增加随机等待。
    """

    def __init__(self, settings: Settings):
        self._settings = settings.scraper
        # 新增：从配置读取并发数限制
        self._concurrency_limit = self._settings.concurrency_limit
        # 恢复：使用您指定的原始配置项
        self._min_delay = self._settings.rate_limit_min_delay
        self._max_delay = self._settings.rate_limit_max_delay
        
        self._semaphore = asyncio.Semaphore(self._concurrency_limit)
        logger.info(
            f"频率控制中间件已启用，最大并发数: {self._concurrency_limit}，"
            f"请求延时范围: [{self._min_delay}, {self._max_delay}] 秒"
        )

    async def process_request(self, target: Target) -> Optional[Union[Target, Response]]:
        await self._semaphore.acquire()
        try:
            delay = random.uniform(self._min_delay, self._max_delay)
            logger.debug(f"速率限制：获得信号量许可，延时 {delay:.2f} 秒，目标 URL: {target.url}")
            await asyncio.sleep(delay)
            return target
        except Exception:
            # 如果在 sleep 或其他地方发生异常，确保释放信号量以防死锁
            self._semaphore.release()
            raise

    async def process_response(self, target: Target, response: Response) -> Optional[Union[Target, Response]]:
        self._semaphore.release()
        return response

    async def process_exception(self, target: Target, exception: Exception) -> Optional[Union[Target, Response]]:
        # 信号量在请求处理流程中，若请求成功发出，则一定被获取。
        # 因此在响应或异常处理中，必须释放。
        self._semaphore.release()
        return None 